feat: lower LogicalNot, AllClose and Phase blocks in the native Expr compiler - #192
Merged
Merged
Conversation
…compiler Closes the misc-block cluster on the Expr op-coverage checklist (#188). Three new lower_block clauses, no new opcodes — each routes through primitives the IR already lowers, bit-identical to the Evaluator. - `Nx.Block.LogicalNot` emits the existing `:logical_not` opcode directly with a trailing coerce, mirroring Emily.Backend.native_logical_not/2. - `Nx.Block.AllClose` composes five existing primitives end-to-end, matching Emily.Backend.native_all_close/4 op-for-op: cast both to the merged float type, compute `abs(a - b) <= atol + rtol * abs(b)`, optionally OR with `isnan(a) AND isnan(b)` for `equal_nan: true`, then reduce over every axis via `:all`. The trailing coerce produces the {:u, 8} predicate dtype. - `Nx.Block.Phase` falls through to the block's composed expansion — `atan2(imag(t), real(t))` — using the TopK-style block-parameter seeding (block-local `:parameter` nodes are fresh and would otherwise resolve to the outer function's input slots). All three primitives in the expansion (atan2/imag/real) are now lowered, so the result is bit-identical to the Evaluator. Probe drops from 12 → 9 misses. Tests cover logical_not (assertion on {:u, 8} dtype), all_close on close/far inputs, custom rtol/atol, and the equal_nan branch, and phase across the four complex quadrants plus the real/imaginary axes.
37 tasks
This was referenced Jun 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes the misc Nx.Block cluster on #188. Three new
lower_blockclauses, no new opcodes — each routes through primitives the IR already lowers, bit-identical to the Evaluator.Nx.Block.LogicalNot— emits the existing:logical_notopcode directly with a trailing coerce, mirroringEmily.Backend.native_logical_not/2.Nx.Block.AllClose— composes five existing primitives end-to-end, matchingEmily.Backend.native_all_close/4op-for-op: cast both to the merged float type, computeabs(a - b) <= atol + rtol * abs(b), optionally OR withisnan(a) AND isnan(b)forequal_nan: true, then reduce over every axis via:all. The trailing coerce produces{:u, 8}.Nx.Block.Phase— falls through to the block's composed expansion (atan2(imag(t), real(t))) via the TopK-style block-parameter seeding (block-local:parameternodes are fresh and would otherwise resolve to the outer function's input slots). All three primitives in the expansion are lowered as of feat: lower 19 more unary ops in the native Expr compiler #189 / feat: lower atan2, quotient and logical_xor in the native Expr compiler #190, so the result is bit-identical to the Evaluator.Probe (
scripts/expr_op_coverage.exs) drops 12 → 9 misses; the misc Nx.Block section of #188 is now empty.Test plan
mix precommitclean (718 tests, 79 properties, 40 doctests, 0 failures — 6 new equivalence tests).compiler_equivalence_test.exs:logical_noton both integer (s32) and float (f32) inputs, asserting the{:u, 8}predicate dtype.all_closeon close inputs (default tolerance → 1), far inputs (→ 0), custom rtol/atol inputs covering boundary cases (atol = 0.2 → close, atol = 0.05 → far), and theequal_nan: truebranch (NaN compares equal to NaN,falseis the default and returns 0).phaseacross the four complex quadrants plus the real and imaginary axes (phase == 0 for positive real, ±pi for negative real, ±pi/2 for imaginary).kOpcodeCount(115).Follows the same pattern as #189 / #190 / #191.